'''
Created on Aug 9, 2012

@author: philipp
'''
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter, LogLocator
import matplotlib.pyplot as plt
import math as m
import pywt

def shift_linear(a, pts, forward=True):
    """\
    Perform a subpixel shift using bilinear interpolation.
    """
    x0, x1 = pts
    if not forward: x0, x1 = -x0, -x1
    a00 = pshift(a, [np.floor(x0), np.floor(x1)], 'nearest')
    a01 = pshift(a, [np.floor(x0), np.ceil(x1)], 'nearest')
    a10 = pshift(a, [np.ceil(x0), np.floor(x1)], 'nearest')
    a11 = pshift(a, [np.ceil(x0), np.ceil(x1)], 'nearest')
    x0 = x0 % 1
    x1 = x1 % 1
    return a00 * (1 - x0) * (1 - x1) + a01 * (1 - x0) * x1 + a10 * x0 * (1 - x1) + a11 * x0 * x1

def pshift(a, ctr, method='linear'):
    """\
    Shift an array periodically so that ctr becomes the origin.
    
    Available methods are 'fourier', 'nearest' and 'linear'. Fourier used to be the default
    but it produced artifacts at strong edges.
    """
    sh = np.array(a.shape)
    ctr = np.asarray(ctr) % sh

    if method.lower() == 'nearest':
        ctri = np.round(ctr).astype(int)
        out = np.empty_like(a)
        c2 = sh - ctri
        if out.ndim == 1:
            out[0:c2[0]] = a[ctri[0]:]
            out[c2[0]:] = a[0:ctri[0]]
    
        elif out.ndim == 2:
            out[0:c2[0], 0:c2[1]] = a[ctri[0]:, ctri[1]:]
            out[0:c2[0], c2[1]:] = a[ctri[0]:, 0:ctri[1]]
            out[c2[0]:, 0:c2[1]] = a[0:ctri[0], ctri[1]:]
            out[c2[0]:, c2[1]:] = a[0:ctri[0], 0:ctri[1]]
        else:
            out[0:c2[0], 0:c2[1], 0:c2[2]] = a[ctri[0]:, ctri[1]:, ctri[2]:]
            out[c2[0]:, 0:c2[1], 0:c2[2]] = a[0:ctri[0], ctri[1]:, ctri[2]:]
            out[0:c2[0], c2[1]:, 0:c2[2]] = a[ctri[0]:, 0:ctri[1], ctri[2]:]
            out[c2[0]:, c2[1]:, 0:c2[2]] = a[0:ctri[0], 0:ctri[1], ctri[2]:]
            out[0:c2[0], 0:c2[1], c2[2]:] = a[ctri[0]:, ctri[1]:, 0:ctri[2]]
            out[c2[0]:, 0:c2[1], c2[2]:] = a[0:ctri[0], ctri[1]:, 0:ctri[2]]
            out[0:c2[0], c2[1]:, c2[2]:] = a[ctri[0]:, 0:ctri[1], 0:ctri[2]]
            out[c2[0]:, c2[1]:, c2[2]:] = a[0:ctri[0], 0:ctri[1], 0:ctri[2]]

    elif method.lower() == 'fourier':
        fout = np.fft.fftn(a.astype(complex))
        out = np.fft.ifftn(fout * np.exp(2j * np.pi * np.sum(fgrid(sh, ctr / sh), axis=0)))

    elif method.lower() == 'linear':
        ctri = np.floor(ctr).astype(int)
        if a.ndim == 1:
            x = ctr - ctri
            a0 = pshift(a, ctri, method='nearest')
            a1 = pshift(a, ctri + 1, method='nearest')
            out = a0 * (1 - x) + a1 * x
        elif a.ndim == 2:   
            x, y = ctr - ctri
            a00 = pshift(a, ctri, method='nearest')
            a01 = pshift(a, ctri + [0, 1], method='nearest')
            a10 = pshift(a, ctri + [1, 0], method='nearest')
            a11 = pshift(a, ctri + [1, 1], method='nearest')
            out = a00 * (1 - x) * (1 - y) + a01 * (1 - x) * y + a10 * x * (1 - y) + a11 * x * y

    return out

def plot3d(arr, name, nr):
    x0 = np.arange(0 , arr.shape[0] , 1.0)
    y0 = np.arange(0 , arr.shape[1] , 1.0)                        
    X0, Y0 = np.meshgrid(x0, y0)
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X0, Y0, arr, rstride=1, cstride=1, cmap=cm.jet, linewidth=0, antialiased=False)
    ax.set_zlim(arr.min(), arr.max()) 
    ax.view_init(elev=77, azim=1)
    ax.zaxis.set_major_locator(LogLocator())
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
    fig.colorbar(surf, shrink=0.5, aspect=5)               
    plt.savefig(name + '_%d.png' % (nr))

def cartesianToPolar(img, r_step=8, r_max=512, phi_steps=256):
    n = img.shape[0]
    ret = np.zeros((r_max / r_step, phi_steps))
    center = np.array([n / 2, n / 2]) - 1
    for r in xrange(0, r_max, r_step):
        for i, phi in enumerate(np.linspace(0.0, 2 * m.pi, num=phi_steps)):
            v1 = getXYvecRounded(r, phi) + center
            v2 = getXYvecRounded(r + 1, phi) + center
            v3 = getXYvecRounded(r, phi + 2 * m.pi / phi_steps) + center
            v4 = getXYvecRounded(r + 1, phi + 2 * m.pi / n) + center       
            tmp = (img[v1[0], v1[1]] + img[v2[0], v2[1]] + img[v3[0], v3[1]] + img[v4[0], v4[1]]) / (4 * r ** 4)
            ret[r / r_step, i] = tmp
    return ret

def getXYvec(r, phi):
    return np.array([r * m.cos(phi), r * m.sin(phi)])

def getXYvecRounded(r, phi):
    return np.round(getXYvec(r, phi))

def moveCenterOfMassToCenter(img):
    ''' COM = 1/M * Sum( m_i * r_i )
    '''    
    n = img.shape[0]
    m = img.shape[1]
    M = 0.0
    sum = 0.0
    for i in range(n):
        for j in range(n):
            sum += img[i, j] * np.array([i, j])
            M += img[i, j]
    COM = sum / M
    
    COM -= np.array([n / 2, m / 2])
    
    img = shift_linear(img, COM)
    
    return img
    
def FFTonAngles(img):
 
    s = img.shape
    ret = np.zeros((s[0], s[1] / 2), dtype='float')
    for i in range(0, s[0]):
        ft = np.fft.fft(img[i, :])[0:s[1] / 2]
        x = np.arange(0, s[1] / 2);
#        plt.plot(x, ft.real)
#        plt.savefig('ft_real_%d.png' % (i))
#        plt.plot(x, ft.imag)
#        plt.savefig('ft_imag_%d.png' % (i))
        ret[i, :] = np.absolute(ft)
    return ret    

def DWTonRadius(img):
    s = img.shape    
    
    length = pywt.dwt_coeff_len(s[0], 4, 'zpd')
    ret1 = np.zeros((length, s[1]))
    ret2 = np.zeros((length, s[1]))
    for i in range(0, s[1]):
        cA, cD = pywt.dwt(img[:, i], 'db2')    
#        x = np.arange(0, s[1] / 2);
#        plt.plot(x, ft.real)
#        plt.savefig('ft_real_%d.png' % (i))
#        plt.plot(x, ft.imag)
#        plt.savefig('ft_imag_%d.png' % (i))
        ret1[:, i] = cA
        #ret2[:, i] = cD
    return ret1#, ret2  
    
def introduceHStripe(img, thickness):
    s = img.shape
    img[(s[0] / 2) - 30:(s[0] / 2) + 30, :] = 0
    return img    

def introduceVStripe(img, thickness):
    s = img.shape
    img[:, (s[1] / 2) - 30:(s[1] / 2) + 30] = 0
    return img  
